import math
import numpy as np
import argparse
import yaml
from pathlib import Path
from enum import Enum

########### Definition ############
class FormationType(Enum):
    wingLike = 'wingLike'
    horiLike = 'horiLike'
    vertLike = 'vertLike'
    rectLike = 'rectLike'
    circLike = 'circLike'

########### utils ############
class InitMsgs:
    """ 初始化信息：载具信息、目标信息、障碍信息
    """
    def __init__(self, ports_cfg, scr_cfg, veh_cfg, tar_cfg, obs_cfg, dcs_cfg):
        self.render_ports_cfg = ports_cfg['render_ports']
        self.server_ports_cfg = ports_cfg['server_ports']
        self.decider_ports_cfg = ports_cfg['decider_ports']

        self.screen_init_msg = scr_cfg
        self.vehicle_init_msg = veh_cfg
        self.target_init_msg = tar_cfg
        self.obstacle_init_msg = obs_cfg
        
        self.decision_cfg = dcs_cfg

def process_yaml_cfg(path_to_configs : str = '/root/code/workspace/src/control_policy/control_policy/MVS/Configs/', 
                     filename : str = 'InitConfigs.yaml'):
    ''' load yaml config
    '''
    parser = argparse.ArgumentParser(description="InitConfigs")
    parser.add_argument('--yaml', default=path_to_configs + filename, type=str,
                        help='the path of config file')
    args = parser.parse_args()
    if args.yaml is not None:
        with open(args.yaml) as f:
            load_cfg = yaml.load(f, Loader=yaml.FullLoader)
            ports_init_msg = load_cfg['ports_cfg']
            screen_init_msg = load_cfg['screen_cfg']
            vehicle_init_msg = load_cfg['vehicle_cfg']
            target_init_msg = load_cfg['target_cfg']
            obstacle_init_msg = load_cfg['obstacle_cfg']
            decision_cfg = load_cfg['decision_cfg']
            # load algo config
            algo_name = decision_cfg['general_cfg']['algo_name']
            algo_mod = __import__(f"config", fromlist=['AlgoConfig'])  # dynamic loading of modules
            algo_cfg = algo_mod.AlgoConfig()
            dcs_cfgs = {'formation_cfg': decision_cfg['formation_cfg'],
                        'general_cfg' : decision_cfg['general_cfg'], 
                        'algo_cfg' : algo_cfg}
            # merge config
            # for cfg_type in dcs_cfgs:
            #     if load_cfg[cfg_type] is not None:
            #         for k, v in load_cfg[cfg_type].items():
            #             setattr(dcs_cfgs[cfg_type], k, v)
            
    return InitMsgs(ports_init_msg, screen_init_msg, vehicle_init_msg,target_init_msg,obstacle_init_msg, dcs_cfgs)

def convert_coordinates(cur_pos, targ_cord_orig_pos, rotate_rad_from_orig_to_targ=math.radians( - 90.), is_flip=False):
    """ 坐标转换函数
    """
    rotate_rad = rotate_rad_from_orig_to_targ
    if is_flip:
        # x,y <- x',y'
        rotate_matrix = np.matrix([[math.cos(rotate_rad), - math.sin(rotate_rad)],
                                [math.sin(rotate_rad), math.cos(rotate_rad)]], dtype=float)
        new_pos = np.matmul(rotate_matrix, cur_pos) + targ_cord_orig_pos
    else:
        # x,y -> x',y'
        rotate_matrix = np.matrix([[math.cos(rotate_rad), math.sin(rotate_rad)],
                                [ - math.sin(rotate_rad), math.cos(rotate_rad)]], dtype=float)
        new_pos = np.matmul(rotate_matrix, cur_pos - targ_cord_orig_pos)
    return new_pos

def calc_formation_offset(formation_set : dict, self_id : int, heading_rad : float):
    # coordinator： NED
    x_offset, y_offset = 0., 0.
    if self_id != formation_set['leader_id']:
        id_offset = self_id - formation_set['leader_id']
        if formation_set['formation_type'] == FormationType.wingLike.value:
            # print("IntoWinglike")
            row = math.ceil(id_offset / 2.)
            col = math.pow(-1, id_offset % 2) * row
            x_offset = - col * formation_set['formation_dist'][0] * math.sin(heading_rad) - row * formation_set['formation_dist'][1] * math.cos(heading_rad)
            y_offset =   col * formation_set['formation_dist'][0] * math.cos(heading_rad) - row * formation_set['formation_dist'][1] * math.sin(heading_rad)
        elif formation_set['formation_type'] == FormationType.horiLike.value:
            x_offset = - id_offset * formation_set['formation_dist'][0] * math.sin(heading_rad)
            y_offset =   id_offset * formation_set['formation_dist'][0] * math.cos(heading_rad)
        elif formation_set['formation_type'] == FormationType.vertLike.value:
            x_offset = - id_offset * formation_set['formation_dist'][1] * math.cos(heading_rad)
            y_offset = - id_offset * formation_set['formation_dist'][1] * math.sin(heading_rad)
    else:
        x_offset, y_offset = 0., 0.
    return x_offset, y_offset

def state_format(state_raw : dict, targets_pos : dict):
    state_new = {}
    for key, values in state_raw.items():
        sur_map = {}
        state_new[key] = [a - b for a, b in zip(targets_pos[key], values[0:2])] + values[2:]
        for idx, val in state_raw.items():
            if idx != key:
                sur_map[idx] = [a - b for a, b in zip(values[0:2], val[0:2])]
        # print(sur_map)
        state_new[key].append(sur_map)
    
    return state_new